
!pip install -q git+https://github.com/ziatdinovmax/gpax@v0.0.3
!pip install -U gdown


from warnings import filterwarnings

import numpy as np
import matplotlib.pyplot as plt

import gpax
import jax.numpy as jnp

from sklearn.model_selection import train_test_split

filterwarnings("ignore", category=UserWarning)

gpax.utils.enable_x64()


#@title Load helper functions  🛠️ { form-width: "20%" }
from typing import Tuple

def make_window(imgsrc: np.ndarray, window_size: int,
                xpos: int, ypos: int) -> np.ndarray:
    """
    Returns the portion of the image within the window given the
    image (imgsrc), the x position and the y position
    """
    imgsrc = imgsrc[int(xpos-window_size/2):int(xpos+window_size/2),
                    int(ypos-window_size/2):int(ypos+window_size/2)]
    return imgsrc


def create_training_set(imgdata: np.ndarray, target: np.ndarray,
                        window_size: int) -> Tuple[np.ndarray]:
    """
    Creates arrays with features (local subimages)
    and targets (corresponding spectra) from hyperspectral data
    """
    feature_arr, target_arr = [], []
    pos = []
    s1, s2 = imgdata.shape[:-1]
    for i in range(s1):
        for j in range(s2):
            arr_loc = make_window(imgdata, window_size, i, j)
            if arr_loc.shape[:-1] != (window_size, window_size):
                continue
            feature_arr.append(arr_loc)
            target_arr.append(target[i, j])
            pos.append([i, j])
    return np.array(pos), np.array(feature_arr), np.array(target_arr)


!gdown https://drive.google.com/uc?id=1VdvoTJ1OaDo-qEIdgZJ1030aK5mBFcQh

Downloading...
From: https://drive.google.com/uc?id=1VdvoTJ1OaDo-qEIdgZJ1030aK5mBFcQh
To: /notebooks/PTO_BEPS_3um.npz
100%|██████████████████████████████████████| 62.1M/62.1M [00:02<00:00, 23.3MB/s]


BEPS = np.load("PTO_BEPS_3um.npz")
BEPS.files

['Image_Amplitude',
 'Image_Phase',
 'Image_Frequency',
 'Image_Q_Factor',
 'Image_Off_Field_Loop_Area',
 'Image_On_Field_Loop_Area',
 'Image_Loop_Width',
 'Image_Average_Coercive_Field',
 'Image_Negative_Coercive_Field',
 'Image_Positive_Coercive_Field',
 'Spectra_DC_Voltage',
 'Spectra_Off_Field_Loop',
 'Spectra_On_Field_Loop',
 'Spectra_On_Field_Frequency_Loop',
 'Spectra_Off_Field_Frequency_Loop']


# DKL inputs
amplitude = BEPS['Image_Amplitude'].astype(np.float64)
phase = BEPS["Image_Phase"].astype(np.float64)
frequency = BEPS["Image_Frequency"].astype(np.float64)
qfactor = BEPS["Image_Q_Factor"].astype(np.float64)
# DKL targets
loop_area = BEPS['Image_Off_Field_Loop_Area']

norm_ = lambda x: (x - x.min()) / x.ptp()

# Concatenate input channels
all_channels = np.concatenate(
    [
     norm_(amplitude)[..., None],
     norm_(phase)[..., None],
     norm_(frequency)[..., None],
     norm_(qfactor)[..., None]
     ], 
     axis=-1)
print(all_channels.shape, loop_area.shape)

# Plot
_, axes = plt.subplots(1, 4, figsize=(20, 5))
for i, ax in enumerate(axes):
    ax.imshow(all_channels[..., i], origin='lower')
    ax.axis('off')

(100, 100, 4) (100, 100)


indices_all, features_all, targets_all = create_training_set(all_channels, loop_area, 12)

#targets_all -= targets_all.mean()
#targets_all /= targets_all.std()

targets_all -= targets_all.min()
targets_all /= targets_all.ptp()


k = 450

_, axes = plt.subplots(1, 4, figsize=(20, 5))
for i, ax in enumerate(axes):
    ax.imshow(features_all[k ,..., i], origin='lower',)
    ax.imshow(features_all[k ,..., i], origin='lower')


def step_mcmc(X, y, X_new):
    key1, key2 = gpax.utils.get_keys()
    data_dim = X.shape[-1]
    dkl = gpax.DKL(data_dim, z_dim=2, kernel='RBF')
    dkl.fit(key1, X, y, num_warmup=333, num_samples=333, num_chains=3, chain_method='vectorized')
    mean, samples = dkl.predict_in_batches(key2, X_new)
    var = samples.squeeze().var(0)
    return mean, var


def get_best_channel(record):
    return record[:,1].argmax()


def update_record(record, action, r):
    new_r = (record[action, 0] * record[action, 1] + r) / (record[action, 0] + 1)
    record[action, 0] += 1
    record[action, 1] = new_r
    return record


def get_reward(obj_history, obj):
    """A reward of +/-1 is given if the integral uncertainty at the current step
    is smaller/larger than the integral uncertainty at the previous step"""
    if jnp.nansum(obj) < obj_history[-1]:
        r = 1
    else:
        r = -1
    return r


def update_datapoints(next_point_idx, train, test):
    X_train, y_train, indices_train = train
    X_test, y_test, indices_test = test
    X_train = jnp.append(X_train, X_test[:, next_point_idx:next_point_idx+1], axis=1)
    X_test = jnp.delete(X_test, next_point_idx, axis=1)
    y_train = jnp.append(y_train, y_test[:, next_point_idx:next_point_idx+1], axis=-1)
    y_test = jnp.delete(y_test, next_point_idx, axis=-1)
    indices_train = jnp.append(indices_train, indices_test[next_point_idx:next_point_idx+1], axis=0)
    indices_test = jnp.delete(indices_test, next_point_idx, axis=0)
    return (X_train, y_train, indices_train), (X_test, y_test, indices_test)


n, h, w, c = features_all.shape
features_all = features_all.reshape(n, h*w, c)

(X_train, X_unmeasured, y_train, y_unmeasured,
 indices_train, indices_unmeasured) = train_test_split(
    features_all, targets_all, indices_all, test_size=0.99, random_state=1)

X_train = X_train.transpose(2, 0, 1)
X_unmeasured = X_unmeasured.transpose(2, 0, 1)
y_train = y_train[None].repeat(X_train.shape[0], axis=0)
y_unmeasured = y_unmeasured[None].repeat(X_train.shape[0], axis=0)

X_train.shape, X_unmeasured.shape, y_train.shape, y_unmeasured.shape

((4, 79, 144), (4, 7842, 144), (4, 79), (4, 7842))


warmup_steps = 5
exploration_steps = 30
# epsilon in epsilon-greedy policy
#eps = np.ones(exploration_steps) * 0.4  # constant
eps = np.linspace(0.4, 0.1, exploration_steps)  # annealing

np.random.seed(55)  # rng seed for epsilon-greedy sampling
record = np.zeros((X_train.shape[0], 2))
channel_choices = []
var_history = []


def plot_acq():
    p = indices_unmeasured[next_point_idx]
    plt.figure(figsize=(6, 6))
    plt.scatter(indices_unmeasured[:, 1], indices_unmeasured[:, 0], c=obj, s=10, marker='s')
    plt.scatter(p[1], p[0], marker='x', s=64, c='red')
    plt.show()
    
# Warm-up phase
for w in range(warmup_steps):
    print("Warmup step {}/{}".format(w+1, warmup_steps))
    # Get/update DKL posterior
    mean, var = step_mcmc(X_train, y_train, X_unmeasured)
    record[:, 0] += 1
    # Select channel that resulted in lowest predictive uncertainty
    var_reduced = np.nansum(var, axis=-1)
    idx = np.argmin(var_reduced)
    # Update records
    var_history.append(var_reduced[idx].item())
    channel_choices.append(idx)
    record[idx, 1] += 1
    # compute the acquisition function for the selected channel
    #obj = mean[idx] + 2*var[idx]
    obj = var[idx]
    # get the next measurement point
    next_point_idx = obj.argmax()
    plot_acq()
    # Do "measurement"
    measured_point = y_unmeasured[:, next_point_idx]
    # update datapoints
    (X_train, y_train, indices_train), (X_unmeasured, y_unmeasured, indices_unmeasured) = update_datapoints(
        next_point_idx, (X_train, y_train, indices_train), (X_unmeasured, y_unmeasured, indices_unmeasured)
        )

# Average over the number of warmup steps
record[:, 1] = record[:, 1] / warmup_steps

# Exploration-exploitation phase
for e in range(exploration_steps):
    print("\nExploration step {}/{}".format(e+1, exploration_steps))
    # Choose channel according to epsilon-greedy policy
    eps_i = eps[e] if e < exploration_steps else eps[-1]
    if np.random.random() > eps_i: 
        idx = get_best_channel(record)
    else:
        idx = np.random.randint(record.shape[0])
    channel_choices.append(idx)
    print("Using channel {}".format(idx+1))
    # Update DKL posterior using the sampled channel
    mean, var = step_mcmc(X_train[idx], y_train[idx], X_unmeasured[idx])
    # Get reward
    r = get_reward(var_history, var)
    # Update records
    record = update_record(record, idx, r)
    var_history.append(jnp.nansum(var).item())
    # Compute acqusition function and do measurement in the suggested point
    #obj = mean + 2*var
    obj = var
    next_point_idx = obj.argmax()
    plot_acq()
    measured_point = y_unmeasured[:, next_point_idx]
    # update datapoints
    (X_train, y_train, indices_train), (X_unmeasured, y_unmeasured, indices_unmeasured) = update_datapoints(
        next_point_idx, (X_train, y_train, indices_train), (X_unmeasured, y_unmeasured, indices_unmeasured)
        )
    # Print the running rewards
    print("\nRewards:")
    for i, r in enumerate(record):
        print("channel {}:  counts {}  reward (avg) {}".format(i+1, (int(r[0])), np.round(r[1], 3)))

    # Store current points
    np.savez("hmcDKL_channels_polar_norm01.npz",
         X_train=X_train, y_train=y_train, X_unmeasured=X_unmeasured,
         y_unmeasured=y_unmeasured, indices_unmeasured=indices_unmeasured,
         indices_measured=indices_train, record=record, channels_id=channel_choices,
         var_history=var_history)

Warmup step 1/5

sample: 100%|██████████| 666/666 [20:36<00:00,  1.86s/it]

                   mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0,0]      7.65      6.66      5.20      1.12     15.83    124.52      1.01
k_length[0,1]      7.41      6.20      5.29      1.09     15.20    143.70      1.01
k_length[1,0]      6.46      6.00      4.37      1.30     13.37    177.93      1.01
k_length[1,1]      6.37      5.60      4.46      0.93     13.08    179.68      1.00
k_length[2,0]      6.01      5.38      4.41      0.57     12.00    146.79      1.01
k_length[2,1]      4.24      4.60      2.57      0.49      9.67    133.14      1.05
k_length[3,0]      5.12      5.29      3.32      0.60     10.30    339.00      1.01
k_length[3,1]      4.71      4.36      3.29      0.70      9.87    299.90      1.00
   k_scale[0]      0.21      0.21      0.15      0.04      0.39    545.33      1.00
   k_scale[1]      0.25      0.19      0.18      0.05      0.49    431.77      1.00
   k_scale[2]      0.28      0.26      0.20      0.04      0.55    555.78      1.00
   k_scale[3]      0.24      0.27      0.18      0.04      0.48    621.41      1.01
     noise[0]      0.01      0.01      0.01      0.00      0.02    200.87      1.00
     noise[1]      0.01      0.00      0.01      0.00      0.01    238.37      1.01
     noise[2]      0.02      0.02      0.01      0.00      0.06     18.80      1.16
     noise[3]      0.02      0.02      0.02      0.00      0.05     68.18      1.01

Warmup step 2/5

sample: 100%|██████████| 666/666 [20:16<00:00,  1.83s/it]

                   mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0,0]      7.01      6.38      4.81      1.06     14.28     68.24      1.04
k_length[0,1]      8.36      7.45      5.94      1.59     16.67    104.35      1.02
k_length[1,0]      5.87      5.54      4.10      1.10     11.48    142.60      1.03
k_length[1,1]      6.28      5.19      4.52      0.95     12.46    146.18      1.02
k_length[2,0]      5.76      6.52      3.78      0.61     12.56    212.28      1.01
k_length[2,1]      5.03      4.75      3.23      0.45     11.11    116.87      1.01
k_length[3,0]      4.73      4.89      2.82      0.69     10.69     80.24      1.02
k_length[3,1]      5.36      5.96      3.39      0.29     12.41     87.49      1.02
   k_scale[0]      0.21      0.21      0.16      0.04      0.37    494.01      1.00
   k_scale[1]      0.24      0.21      0.18      0.04      0.46    577.60      1.00
   k_scale[2]      0.31      0.26      0.23      0.03      0.67    302.70      1.01
   k_scale[3]      0.23      0.23      0.16      0.04      0.43    366.60      1.01
     noise[0]      0.01      0.01      0.01      0.00      0.02    210.32      1.00
     noise[1]      0.01      0.00      0.01      0.00      0.01    234.36      1.01
     noise[2]      0.02      0.02      0.02      0.00      0.05     18.77      1.14
     noise[3]      0.02      0.01      0.01      0.00      0.04     53.75      1.03

Warmup step 3/5

sample: 100%|██████████| 666/666 [20:59<00:00,  1.89s/it]

                   mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0,0]      7.51      6.94      5.05      0.98     16.15     66.18      1.04
k_length[0,1]      8.43      8.30      5.68      0.97     17.50     95.33      1.03
k_length[1,0]      6.79      6.49      4.62      1.29     13.79    114.86      1.00
k_length[1,1]      6.29      5.08      4.74      1.03     11.93    129.11      1.01
k_length[2,0]      4.94      5.01      3.03      0.26     11.23     81.28      1.01
k_length[2,1]      5.89      6.09      3.81      0.41     12.46     71.91      1.02
k_length[3,0]      5.42      5.26      3.52      0.53     11.96    178.48      1.00
k_length[3,1]      4.88      5.25      2.91      0.51     11.64     89.41      1.01
   k_scale[0]      0.20      0.18      0.15      0.04      0.39    360.75      1.00
   k_scale[1]      0.26      0.24      0.19      0.03      0.49    498.69      1.00
   k_scale[2]      0.30      0.28      0.21      0.03      0.61    220.40      1.02
   k_scale[3]      0.26      0.28      0.19      0.04      0.47    370.66      1.01
     noise[0]      0.01      0.01      0.01      0.00      0.02    127.84      1.01
     noise[1]      0.01      0.00      0.01      0.00      0.01    206.99      1.00
     noise[2]      0.02      0.02      0.01      0.00      0.05     41.09      1.07
     noise[3]      0.02      0.01      0.01      0.00      0.03     18.85      1.17

Warmup step 4/5


dataset = np.load("hmcDKL_channels_polar_norm01.npz")

X_train=dataset["X_train"]
y_train=dataset["y_train"]
X_unmeasured=dataset["X_unmeasured"]
y_unmeasured=dataset["y_unmeasured"]
indices_unmeasured=dataset["indices_unmeasured"]
indices_train=dataset["indices_measured"]
record=dataset["record"]
channel_choices=dataset["channels_id"].tolist()
var_history=dataset["var_history"].tolist()


for i, r in enumerate(record):
    print("channel {}:  counts {}  reward (avg) {}".format(i+1, (int(r[0])), np.round(r[1], 3)))

channel 1:  counts 7  reward (avg) -0.286
channel 2:  counts 29  reward (avg) 0.448
channel 3:  counts 7  reward (avg) -0.286
channel 4:  counts 7  reward (avg) -0.286


plt.figure(figsize = (5, 3), dpi=100)
plt.plot(np.arange(1, exploration_steps+warmup_steps+1), var_history, c='k')
plt.scatter(np.arange(1, exploration_steps+warmup_steps+1), var_history,
            c=[m+1 for m in channel_choices],
            s=128, cmap='rainbow', alpha=1)
plt.xlabel("Exploration step", fontsize=14)
plt.ylabel("Uncertainty", fontsize=14)
plt.xticks(size=14)
plt.yticks(size=14)
cbar = plt.colorbar(label = "Channel Number")
cbar.set_ticks([1, 2, 3, 4])


mean, var = step_mcmc(X_train, y_train, X_unmeasured)

sample: 100%|██████████| 666/666 [29:42<00:00,  2.68s/it]

                   mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0,0]      6.04      4.32      4.79      2.05     10.29     40.34      1.06
k_length[0,1]     11.93      8.18      9.88      2.19     22.15     51.32      1.07
k_length[1,0]      5.70      5.26      4.12      1.54     10.31     53.43      1.04
k_length[1,1]     10.51      8.96      8.53      1.77     20.10    108.32      1.01
k_length[2,0]      8.44     11.94      5.60      0.83     16.22     51.53      1.08
k_length[2,1]      3.80      3.64      2.69      0.77      7.44     54.60      1.08
k_length[3,0]      4.33      5.02      2.50      0.70      9.84     41.81      1.06
k_length[3,1]      8.29      7.40      6.75      0.86     15.97     41.88      1.05
   k_scale[0]      0.21      0.18      0.16      0.04      0.38    176.10      1.01
   k_scale[1]      0.23      0.19      0.18      0.05      0.45    332.78      1.02
   k_scale[2]      0.38      0.29      0.28      0.06      0.77    129.45      1.02
   k_scale[3]      0.30      0.27      0.23      0.05      0.58    106.15      1.05
     noise[0]      0.01      0.00      0.01      0.00      0.01     99.56      1.01
     noise[1]      0.00      0.00      0.00      0.00      0.01     56.03      1.07
     noise[2]      0.00      0.00      0.00      0.00      0.01     85.43      1.04
     noise[3]      0.00      0.00      0.00      0.00      0.01     16.82      1.10


from mpl_toolkits.axes_grid1 import make_axes_locatable

for m, v in zip(mean, var):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 6))
    im1 = ax1.scatter(indices_unmeasured[:, 1], indices_unmeasured[:, 0], c=m, s=10, marker='s')
    ax1.set_aspect('equal')
    divider = make_axes_locatable(ax1)
    cax = divider.append_axes('right', size='3%', pad=0.1)
    fig.colorbar(im1, cax=cax, orientation='vertical')
    im2 = ax2.scatter(indices_unmeasured[:, 1], indices_unmeasured[:, 0], c=v, s=10, marker='s')
    ax2.set_aspect('equal')
    divider = make_axes_locatable(ax2)
    cax = divider.append_axes('right', size='3%', pad=0.1)
    fig.colorbar(im2, cax=cax, orientation='vertical')
    plt.show()


print(jnp.sum((abs(mean - y_unmeasured)), axis=-1))

[1184.67279246 1005.68638224 1411.77420842 1554.34109267]

Active Channel Learning with Fully Bayesian Deep Kernel Learning¶

Import and Download data¶